from mpi4py import MPI
import numpy as np
import sys


def matrix_multiply(size, blocks):
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size_sqrt = int(np.sqrt(comm.Get_size()))

    assert size % size_sqrt == 0

    block_size = size // size_sqrt

    if rank == 0:
        matrix_a = np.random.rand(size, size)
        matrix_b = np.random.rand(size, size)
    else:
        matrix_a = None
        matrix_b = None

    matrix_a = comm.bcast(matrix_a, root=0)
    matrix_b = comm.bcast(matrix_b, root=0)

    local_matrix_a = matrix_a[rank // size_sqrt * block_size:(rank // size_sqrt + 1) * block_size, :]
    local_matrix_b = matrix_b[:, rank % size_sqrt * block_size:(rank % size_sqrt + 1) * block_size]

    local_result = np.dot(local_matrix_a, local_matrix_b)

    result = comm.gather(local_result, root=0)

    if rank == 0:
        result = np.concatenate(result)
        return result


if __name__ == '__main__':
    size = int(sys.argv[1])
    blocks = int(sys.argv[2])
    result = matrix_multiply(size, blocks)
    if MPI.COMM_WORLD.Get_rank() == 0:
        print("Result:", result)

